import pandas as pd

def process_csv(input_file, output_file):
    # Load Data
    df = pd.read_csv(input_file)

    # Remove repeated headers and excess blank rows
    df = df.drop_duplicates(keep='first')
    df = df.dropna(how='all')

    # Standardize Headers
    column_mapping = {
        '序号': 'serial_no',
        '学校名称': 'school_name',
        '学校性质': 'school_type',
        '教职工人数': 'total_staff_count',
        '专职教师人数': 'full_time_teacher_count',
        '师生比': 'student_faculty_ratio_raw'
    }
    df = df.rename(columns=column_mapping)

    # Filter Data
    df = df[df['school_type'] == '公办']

    # Process Ratio
    df = df[df['student_faculty_ratio_raw'].str.match(r'^1:\d+$')]
    df['student_faculty_ratio'] = df['student_faculty_ratio_raw'].str.extract(r'^1:(\d+)$').astype(int)

    # Calculate Population Metrics
    df['estimated_student_count'] = df['total_staff_count'] * df['student_faculty_ratio']
    df['students_per_ft_teacher'] = df.apply(lambda row: row['estimated_student_count'] / row['full_time_teacher_count'] if row['full_time_teacher_count'] != 0 else 0, axis=1)

    # Ensure Numeric Data
    numeric_columns = ['total_staff_count', 'full_time_teacher_count', 'estimated_student_count', 'students_per_ft_teacher']
    df[numeric_columns] = df[numeric_columns].astype(float)

    # Handle Missing Data
    df = df.dropna(subset=numeric_columns)

    # Select & Order Final Columns
    final_columns = ['serial_no', 'school_name', 'total_staff_count', 'full_time_teacher_count', 'estimated_student_count', 'students_per_ft_teacher']
    df = df[final_columns]

    # Save Output
    df.to_csv(output_file, index=False)

# Process each file separately
input_files = ['baoshan-schools-2025.csv', 'baoshan-schools-2024.csv', 'baoshan-schools-2020.csv']
output_files = ['clean-baoshan-schools-2025.csv', 'clean-baoshan-schools-2024.csv', 'clean-baoshan-schools-2020.csv']

for input_file, output_file in zip(input_files, output_files):
    process_csv(input_file, output_file)